# -*- coding: utf-8 -*-
__author__ = 'BuddyHolly'

import cv2
import random
import numpy as np

ALPHA = 3
SEED = 5949


# ===================================
# *** 图片水印提取 ***
#
# img_path_name：原图路径
# ori_path_name：添加过水印图片路径
# res_path_name：提取水印保存路径
# alpha：通道
# ran_seed：随机所使用seed
#
# 提取水印图存在res_path_name处，无返回值
# ===================================

def de_pic(img_path_name, ori_path_name, res_path_name, alpha=ALPHA, ran_seed=SEED):
    # 输入原图与添加过水印的图
    ori = cv2.imread(ori_path_name)
    img = cv2.imread(img_path_name)
    # 原图快速傅里叶变换
    ori_f = np.fft.fft2(ori)

    # 如果图像尺寸不同则返回错误
    if ori.shape[0] < img.shape[0] or ori.shape[1] < img.shape[1]:
        return True

    # 计算图像大小
    height, width = ori.shape[0], ori.shape[1]

    # 补全水印图
    tmp = np.zeros(ori.shape)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            tmp[i][j] = img[i][j]

    # 水印图快速傅里叶变换
    img_f = np.fft.fft2(tmp)

    # 图像相减得到水印
    watermark = (img_f - ori_f) / alpha
    watermark = np.real(watermark)
    res = np.zeros(watermark.shape)
    x, y = range(height / 2), range(width)

    # 水印图像解密
    random.seed(height + width + ran_seed)
    random.shuffle(x)
    random.shuffle(y)

    # 恢复水印图
    for i in range(height / 2):
        for j in range(width):
            res[x[i]][y[j]] = watermark[i][j]

    cv2.imwrite(res_path_name, res, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
